/***************************************************************************
 *   Copyright (C) 2009 by Zheng Yonglei
 *   zhengyonglei@gmail.com   *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <QString>
#include <QTextStream>

#include "SQLiteDB.h"
#include "TrainerThread.h"
#include "Options.h"
#include "Logger.h"

TrainerThread::TrainerThread(QObject * parent) :
	QThread(parent),
	mIsTraining(false),
	mTrainData(NULL),
	mAnn(NULL)
{
	QObject::connect(&mTrainTimer, SIGNAL(timeout()), this, SLOT(trainNextEpoch()));

	mAnnFile = QString(USER_DIR) + "walter" + QString("/ann.txt");

	mInputNameTestCount = QString("test_count");
	mInputNameIntv = QString("test_interval");
	mInputNameScore = QString("test_record_history_score");
	mInputNameGoogleCount = QString("test_word_self_factor");

	mInputNameList<<mInputNameTestCount;
	mInputNameList<<mInputNameIntv;
	mInputNameList<<mInputNameScore;
	mInputNameList<<mInputNameGoogleCount;
}

TrainerThread::~TrainerThread()
{
}

void TrainerThread::run()
{
	LOG_REF.setThreadID("TrainerThread");
	LOG<<"Trainer thread is running."<<endl;
	mAnnFile = QString(USER_DIR) + DB.curUsr + QString("/ann.txt");
	startTraining();
	mEpochMS = OPT.trainEpochMS;
	mCurEpoch = 0;
	mTrainTimer.setSingleShot(true);
	mTrainTimer.start(mEpochMS);
}

void TrainerThread::exit(int returnCode)
{
	if (mIsTraining)
		stopTraining();
	QThread::exit(returnCode);
}

void TrainerThread::startTraining()
{
	if (mTrainData)
		fann_destroy_train(mTrainData);

	mTrainData = fann_read_train_from_file(mTrainDataFile.toUtf8().data());

	if (QFile::exists(mAnnFile))
		mAnn = fann_create_from_file(mAnnFile.toUtf8().data());
	else
	{
		//create ANN
		const unsigned int num_input = INPUT_SIZE;
		const unsigned int num_output = 1;
		const unsigned int num_layers = 5;
		//G-->7  R-->18 E-->5
		const unsigned int num_layer1_neurons_hidden = 7;
		const unsigned int num_layer2_neurons_hidden = 18;
		const unsigned int num_layer3_neurons_hidden = 5;

		mAnn = fann_create_standard(num_layers,
				num_input,
				num_layer1_neurons_hidden,
				num_layer2_neurons_hidden,
				num_layer3_neurons_hidden,
				//num_layer4_neurons_hidden,
				num_output);

		fann_set_activation_function_hidden(mAnn, FANN_SIGMOID_SYMMETRIC);
		fann_set_activation_function_output(mAnn, FANN_SIGMOID);

		//The output should be extreme (0 for fail or 1 for pass)
		//fann_set_activation_steepness_output(mAnn, 1.0f);
		fann_set_activation_steepness_output(mAnn, 2.0f);
	}

	//scale the input and output data
	fann_set_scaling_params(
				mAnn,
				mTrainData,
				0,	/* New input minimum */
				1,	/* New input maximum */
				0,	/* New output minimum */
				1);	/* New output maximum */
	fann_scale_train( mAnn, mTrainData );

	mIsTraining = true;

	emit statusChanged("Training...");
}

void TrainerThread::setTrainData()
{
	mTrainDataFile = QString(USER_DIR) + DB.curUsr + QString("/train_data.txt");

	QString command;
	QTextStream cs(&command);

	cs<<	" select "
			" test_interval, "
			" test_count, "
			" test_record_history_score, "
			" test_word_self_factor, "
			" test_result "
			" from test_train_data ";
	DB.quickGet(command);

	QFile file(mTrainDataFile);
	if (!file.open(QIODevice::WriteOnly | QIODevice::Text | QIODevice::Truncate))
		return;

	QTextStream out(&file);
	clearTestDataVect();
	int dataCount = DB.getRowCount();
	//out<<dataCount<<" 4 1"<<endl;
	out<<dataCount<<" "<<INPUT_SIZE<<" 1"<<endl;
	for (int i=0; i<DB.getRowCount(); i++)
	{
		QStringList inputValList;

		fann_type* input = new fann_type[INPUT_SIZE]; //max input

		if (mInputNameList.contains(mInputNameTestCount))
		{
			QString testCountValStr = DB.getTableValue(i, mInputNameTestCount);
			input[mInputNameList.indexOf(mInputNameTestCount)] = testCountValStr.toFloat();
			inputValList<<testCountValStr;
		}

		if (mInputNameList.contains(mInputNameIntv))
		{
			QString testIntervalValStr = DB.getTableValue(i, mInputNameIntv);
			input[mInputNameList.indexOf(mInputNameIntv)] = testIntervalValStr.toFloat();
			inputValList<<testIntervalValStr;
		}

		if (mInputNameList.contains(mInputNameScore))
		{
			QString testScoreValStr = DB.getTableValue(i, mInputNameScore);
			input[mInputNameList.indexOf(mInputNameScore)] = testScoreValStr.toFloat();
			inputValList<<testScoreValStr;
		}

		if (mInputNameList.contains(mInputNameGoogleCount))
		{
			QString googleSearchCountValStr = DB.getTableValue(i, mInputNameGoogleCount);
			input[mInputNameList.indexOf(mInputNameGoogleCount)] = googleSearchCountValStr.toInt();
			inputValList<<googleSearchCountValStr;
		}

		mTestInputVect.push_back(input);

		for (int j=0; j<inputValList.size(); j++)
			out<<inputValList[j]<<" ";
		out<<endl;
		//out<<testCount<<" "<<testInterval<<" "<<testScore<<" "<<googleSearchCount<<endl;

		QString testResult = DB.getTableValue(i, "test_result");
		mTestOutVect.push_back(testResult.toInt());
		out<<testResult<<endl;
	}
}

void TrainerThread::clearTestDataVect()
{
	for (int i=0; i<mTestInputVect.size(); i++)
		delete[] mTestInputVect[i];

	mTestInputVect.clear();
	mTestOutVect.clear();
}

void TrainerThread::trainNextEpoch()
{
	if (mIsTraining)
	{
		mCurEpoch++;
		mFinalMSE = fann_train_epoch(mAnn, mTrainData);
		mMSEVect.push_back(mFinalMSE);
		if (mCurEpoch <= 1)
			mInitMSE = mFinalMSE;

		emit updateMSE(mCurEpoch, mFinalMSE);

		if (checkTrainingFinished())
			stopTraining();
		else
			mTrainTimer.start(mEpochMS);
	}
}

void TrainerThread::stopTraining()
{
	LOG<<"Stop trainer thread"<<endl;
	QString annTmp = mAnnFile + QString(".tmp");
	fann_save(mAnn, annTmp.toUtf8().data());
	if (mFinalMSE < mInitMSE)
	{
		LOG<<"mFinalMSE("<<mFinalMSE<<") < mInitMSE("<<mInitMSE<<") --> Save ANN"<<endl;
		if (fann_save(mAnn, mAnnFile.toUtf8().data()) == -1)
			LOG<<"Save ANN file failed. mAnnFile="<<mAnnFile.toUtf8().data()<<endl;
	}
	else
	{
		LOG<<"mFinalMSE("<<mFinalMSE<<") >= mInitMSE("<<mInitMSE<<") --> Do NOT save ANN"<<endl;
	}

	fann_destroy_train(mTrainData);
	mTrainData = NULL;
	fann_destroy(mAnn);
	mAnn = NULL;
	mIsTraining = false;

	emit statusChanged("Training stopped...");
}

bool TrainerThread::checkTrainingFinished()
{
	if (mCurEpoch > OPT.trainEpochCountMax)
		return true;

	if (mMSEVect.size() < OPT.checkFinishedEpochCount)
		return false;

	float mean = 0.0f;
	float mse = 0.0f;
	for (	int i = mMSEVect.size() - OPT.checkFinishedEpochCount; i < mMSEVect.size(); i++)
		mean += mMSEVect[i];

	mean /= float(OPT.checkFinishedEpochCount);
	LOG<<"mean="<<mean<<endl;

	for (	int i = mMSEVect.size() - OPT.checkFinishedEpochCount; i < mMSEVect.size(); i++)
	{
		mse += (mean - mMSEVect[i]) * (mean - mMSEVect[i]);
	}

	mse = sqrt(mse/(OPT.checkFinishedEpochCount - 1));
	LOG<<"mse="<<mse<<endl;

	LOG<<"OPT.finishedMSEOfMSEs="<<OPT.finishedMSEOfMSEs<<endl;
	if (mse <= OPT.finishedMSEOfMSEs)
		return true;
	else
		return false;
}

float TrainerThread::testANN(QString annFile)
{
	if (annFile.isEmpty())
		annFile = mAnnFile;

	if (!QFile::exists(annFile))
	{
		LOG<<"Failed to open ANN file: annFile="<<annFile.toUtf8().data()<<endl;
		return 0;
	}

	struct fann *ann = fann_create_from_file(annFile.toUtf8().data());
	float testSimScore[10];
	int testSimCorrectSum[10];
	int testSimSum[10];
	for (int i=0; i<10; i++)
	{
		testSimScore[i] = 0.0f;
		testSimSum[i] = 0;
		testSimCorrectSum[i] = 0;
	}

	int sampleCount = 0;
	int correctCount = 0;
	float testValueSum = 0;
	assert(mTestInputVect.size() == mTestOutVect.size());
	for (int i=0; i<mTestInputVect.size(); i++)
	{
		fann_type input[INPUT_SIZE];
		for (int j=0; j<INPUT_SIZE; j++)
			input[j] = mTestInputVect[i][j];

		fann_scale_input(ann, input);
		fann_type* calc_out = fann_run(ann, input);
		fann_descale_output( ann, calc_out );
		float testValue = calc_out[0];

		//LOG<<"TrainerThread::testANN-->testValue="<<testValue<<endl;
		//LOG<<"TrainerThread::testANN-->correctCount="<<correctCount<<endl;
		int testValClass = int(testValue*10.0f);
		if (testValClass < 0 || testValClass > 10)
		{
			LOG<<"Error! testValue="<<testValue<<endl;
			continue;
		}
		else
		{
			sampleCount++;
			int testValInt = testValue*2;
			if (testValInt == mTestOutVect[i])
			{
				correctCount++;
				testSimCorrectSum[testValClass]++;
			}
		}


		//LOG<<"TrainerThread::testANN-->testValInt="<<testValInt<<endl;
		if (mTestOutVect[i] == 0)
		{
			testValue = 1.0f - testValue;
			//LOG<<"TrainerThread::testANN-->testValue="<<testValue<<endl;
		}

		testValueSum += testValue;

		//LOG<<"TrainerThread::testANN-->testValClass="<<testValClass<<endl;
		testSimScore[testValClass] += testValue;
		testSimSum[testValClass]++;
	}

	QString MFileName(QString("./ann_test_data/train_test_%1.m").arg(QDateTime::currentDateTime().toString("yy_MM_dd_hh_mm")));
	QFile MFile(MFileName);
	if (!MFile.open(QIODevice::WriteOnly | QIODevice::Text))
		return -1;

	QTextStream ms(&MFile);
	ms<<"%inputs:"<<mInputNameList.join("\t")<<endl;
	ms<<"train_epoch="<<mCurEpoch<<";"<<endl;
	ms<<"sample_count="<<sampleCount<<";"<<endl;
	ms<<"test_correct_count="<<correctCount<<";"<<endl;
	ms<<"test_float_result_sum="<<testValueSum<<";"<<endl;

	LOG<<"correctCount="<<correctCount<<endl;
	int testDataCount = mTestInputVect.size();
	//LOG<<"TrainerThread::testANN-->failSample="<<failSample<<endl;
	LOG<<"testValueSum/testDataCount="<<testValueSum/testDataCount<<endl;
	LOG<<"correctCount/testDataCount="<<float(correctCount)/testDataCount<<endl;

	ms<<"%Group\tSample_Count\tScore\tCorrect_Count"<<endl;
	ms<<"test_data=["<<endl;
	for (int i=0; i<10; i++)
	{
		LOG	<<"Group "<<i<<": testSimSum="<<testSimSum[i]
		    	<<"\ttestSimScore="<<float(testSimScore[i])/float(testSimSum[i])
		    	<<"\tcorrectRate="<<float(testSimCorrectSum[i])/float(testSimSum[i])<<endl;


		ms<<i<<"\t"<<testSimSum[i]<<"\t"<<testSimScore[i]<<"\t"<<testSimCorrectSum[i]<<endl;
	}
	ms<<"];"<<endl;

	fann_destroy(ann);

	float accuracy = testValueSum/testDataCount;
	LOG<<"accuracy="<<accuracy<<endl;
	return accuracy;
}

int TrainerThread::calculateScore(const QStringList& recordList)
{
	int testScore = 0;
	for (int i=recordList.size()-1; i >= 0; --i)
	{
		const QString& record = recordList[i];
		QString testResult = record.right(record.size() - record.indexOf(":") - 1);
		if (testResult.isEmpty())
			continue;

		if (i == recordList.size()-1)
		{
			QString lastTestDateStr = record.left(record.indexOf(":"));

			if (testResult[0] == 'Y')
				testScore = 1;
			else if (testResult[0] == 'N')
				testScore = -1;
			else
			{
				LOG<<"Illegal testResult:"<<testResult<<endl;
				assert(0);
			}
		}
		else
		{
			if (testResult[0] == 'Y')
			{
				if (testScore > 0)
					testScore++;
				else
					break;
			}
			else if (testResult[0] == 'N')
			{
				if (testScore < 0)
					testScore--;
				else
					break;
			}
			else
			{
				LOG<<"Illegal testResult:"<<testResult<<endl;
				assert(0);
			}
		}
		LOG<<"testResult: "<<testResult<<endl;
	}

	return testScore;
}

float TrainerThread::simInput(
		fann *ann,
		int testCount,
		int testIntv,
		int testScore,
		int googleSearchCount)
{
	fann_type input[INPUT_SIZE];
	int pos = 0;
	if (mInputNameList.contains(mInputNameTestCount))
	{
		input[pos] = testCount;
		pos++;
	}

	if (mInputNameList.contains(mInputNameIntv))
	{
		input[pos] = testIntv;
		pos++;
	}

	if (mInputNameList.contains(mInputNameScore))
	{
		input[pos] = testScore;
		pos++;
	}

	if (mInputNameList.contains(mInputNameGoogleCount))
	{
		input[pos] = googleSearchCount;
		pos++;
	}

	fann_scale_input(ann, input);
	fann_type* calc_out = fann_run(ann, input);
	fann_descale_output( ann, calc_out );
	float testValue = calc_out[0];

	return testValue;
}

void TrainerThread::simSingleWord(
			QString word,
			QString test_records,
			int googleSearchCount,
			int size,
			fann *ann,
			vector<int>& intvVect,
			vector<float>& simValVect)
{
	QStringList recordList = test_records.split("|");
	int testCount = recordList.size();
	if (recordList.isEmpty())
		return;

	int testScore = calculateScore(recordList);

	//LOG<<"word="<<word<<endl;
	//LOG<<"last test date: "<<lastTestDate.toString(DATE_FORMAT)<<endl;
	//LOG<<"testScore: "<<testScore<<endl;

	int intvMax = size;
	for (int i=1; i <= intvMax; i++)
	{
		int testIntv = i;
		float testValue = simInput(ann, testCount, testIntv, testScore, googleSearchCount);

		intvVect.push_back(testIntv);
		simValVect.push_back(testValue);
		LOG<<"testValue="<<testValue<<endl;
	}
}

void TrainerThread::testWord(int id, int size)
{
	QString command;
	QTextStream cs(&command);
	cs<<" select words.word, words.google_search_count, test_records.test_record "
			" from words,test_records "
			" where "
			" words.id = test_records.id "
			" and words.id = "<<id<<";";
	DB.quickGet(command);
	if (DB.getRowCount() <= 0)
		return;

	QString word = DB.getTableValue(0, "word");
	QString test_records = DB.getTableValue(0, "test_record");
	int googleSearchCount = DB.getTableValue(0, "google_search_count").toFloat();

	QString annFile = mAnnFile;
	if (!QFile::exists(annFile))
	{
		LOG<<"Failed to open ANN file: annFile="<<annFile.toUtf8().data()<<endl;
		return;
	}

	struct fann *ann = fann_create_from_file(annFile.toUtf8().data());
	vector<int> intvVect;
	vector<float> simValVect;

	simSingleWord(word, test_records, googleSearchCount, size, ann, intvVect, simValVect);

	fann_destroy(ann);

	emit updateWordSim(word, intvVect, simValVect);
}

void TrainerThread::simAvg()
{
	QString command;
	QTextStream cs(&command);
	cs<<" select words.word, words.google_search_count, test_records.test_record "
			" from words,test_records "
			" where "
			" words.id = test_records.id; ";
	DB.quickGet(command);

	int googleSearchCountSum = 0;
	for (int i=0; i<DB.getRowCount(); i++)
	{
		int googleSearchCount = DB.getTableValue(0, "google_search_count").toFloat();
		googleSearchCountSum += googleSearchCount;
	}

	int  googleSearchCountAvg = googleSearchCountSum/DB.getRowCount();
	//LOG<<"googleSearchCountAvg="<<googleSearchCountAvg<<endl;

	QString annFile = mAnnFile;
	if (!QFile::exists(annFile))
	{
		LOG<<"Failed to open ANN file: annFile="<<annFile.toUtf8().data()<<endl;
		return;
	}

	struct fann *ann = fann_create_from_file(annFile.toUtf8().data());
	vector<int> intvVect;
	vector<float> simValVect;

	int intvMax = 30;
	int testCountMax = 5;
	int testScore = 0;
	for (int i=1; i <= testCountMax; i++)
	{
		intvVect.clear();
		simValVect.clear();
		for (int j=1; j<=intvMax; j++)
		{
			int testIntv = j;
			int testCount = i;

			float testValue = simInput(ann, testCount, testIntv, testScore, googleSearchCountAvg);

			intvVect.push_back(testIntv);
			simValVect.push_back(testValue);
			//LOG<<"testValue="<<testValue<<endl;
		}

		emit updateWordSim("avg", intvVect, simValVect);
	}

	fann_destroy(ann);
}
